{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Vascular Model Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains an example of how to train a vascular model for automatic vessel segmentation from a CTA. For this example we are using CTA of the head. This example requires an expert mask as an input. Expert mask is a binary image volume, where vessels are marked 1 and rest is 0. It also requires one more mask, which serves as mask of the brain region with in the head CTA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If ITK is not installed in your python environment, you need to define the environment variable `ITK_BUILD_DIR` that contains the path to where ITK was built."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to find the directory in which TubeTK was build. This is required to find the path to the testing data, and may be also required to find the TubeTK library paths if your python environment does not include it.\n",
    "The environment variable `TubeTK_BUILD_DIR` needs to be defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Path for TubeTK libs and bin\n",
    "#Values takend from TubeTK launcher\n",
    "\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/SlicerExecutionModel-build/GenerateCLP/\")\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/SlicerExecutionModel-build/GenerateCLP/Release\")\n",
    "\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/ITK-build/bin/\")\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/ITK-build/bin/Release\")\n",
    "\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/TubeTK-build/bin/\")\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/TubeTK-build/bin/Release\")\n",
    "sys.path.append(\"C:/src/TubeTK_Python_ITK/TubeTK-build/lib/\")\n",
    "sys.path.append(\"C:/src/TubeTK_Python_ITK/TubeTK-build/lib/Release\")\n",
    "\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/VTK-build/bin/\")\n",
    "#sys.path.append(\"C:/src/TubeTK_Python_ITK/VTK-build/bin/Release\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Setting TubeTK Build Directory\n",
    "\n",
    "TubeTK_BUILD_DIR=None\n",
    "if 'TubeTK_BUILD_DIR' in os.environ:\n",
    "    TubeTK_BUILD_DIR = os.environ['TubeTK_BUILD_DIR']\n",
    "else:\n",
    "    print('TubeTK_BUILD_DIR not found!')\n",
    "    print('  Set environment variable')\n",
    "    os.environ[\"TubeTK_BUILD_DIR\"] = \"C:/src/TubeTK_Python_ITK/TubeTK-build\"\n",
    "    TubeTK_BUILD_DIR = os.environ[\"TubeTK_BUILD_DIR\"]\n",
    "    #sys.exit( 1 )\n",
    "    \n",
    "if not os.path.exists(TubeTK_BUILD_DIR):\n",
    "    print('TubeTK_BUILD_DIR set by directory not found!')\n",
    "    print('  TubeTK_BUILD_DIR = ' + TubeTK_BUILD_DIR )\n",
    "    sys.exit(1)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Setting ITK Build Directory and importing ITK\n",
    "\n",
    "try:\n",
    "    import itk\n",
    "except:\n",
    "    ITK_BUILD_DIR = None\n",
    "    if 'ITK_BUILD_DIR' in os.environ:\n",
    "        ITK_BUILD_DIR = os.environ['ITK_BUILD_DIR']\n",
    "    else:\n",
    "        print('ITK_BUILD_DIR not found!')\n",
    "        print('  Set environment variable')\n",
    "        os.environ[\"ITK_BUILD_DIR\"] = \"C:/src/TubeTK_Python_R/ITK-build\"\n",
    "        ITK_BUILD_DIR = os.environ[\"ITK_BUILD_DIR\"]\n",
    "        #sys.exit( 1 )\n",
    "\n",
    "    if not os.path.exists(ITK_BUILD_DIR):\n",
    "        print('ITK_BUILD_DIR set by directory not found!')\n",
    "        print('  ITK_BUIDL_DIR = ' + ITK_BUILD_DIR )\n",
    "        sys.exit(1)\n",
    "    # Append ITK libs\n",
    "    sys.path.append(\"C:/src/TubeTK_Python_ITK/ITK-build/Wrapping/Generators/Python/Release\")\n",
    "    sys.path.append(\"C:/src/TubeTK_Python_ITK/ITK-build/lib/Release\")\n",
    "    sys.path.append(\"C:/src/TubeTK_Python_ITK/ITK-build/lib\")\n",
    "    \n",
    "    # Append TubeTK libs\n",
    "    sys.path.append(\"C:/src/TubeTK_Python_ITK/TubeTK-build/ITKModules/TubeTKITK-build/Wrapping/Generators/Python/Release\")\n",
    "    import itk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from itk import TubeTKITK as itktube\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Dimension = 3\n",
    "PixelType = itk.F\n",
    "\n",
    "CTImageFileName = os.path.join(TubeTK_BUILD_DIR, 'MIDAS_Data\\inputCTA.mha')\n",
    "ExpertMaskImageFileName = os.path.join(TubeTK_BUILD_DIR, 'MIDAS_Data\\inputExpertMask.mha')\n",
    "MaskImageFileName = os.path.join(TubeTK_BUILD_DIR, 'MIDAS_Data\\inputMask.mha')\n",
    "\n",
    "\n",
    "SpatialObjectType = itk.SpatialObject[Dimension]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read the input images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2000D\u001b[KitkImageFileReaderIF3: 0.000000\u001b[2000D\u001b[KitkImageFileReaderIF3: 1.000000\u001b[2000D\u001b[K\u001b[2000D\u001b[KitkImageFileReaderIF3: 0.000000\u001b[2000D\u001b[KitkImageFileReaderIF3: 1.000000\u001b[2000D\u001b[K\u001b[2000D\u001b[KitkImageFileReaderIF3: 0.000000\u001b[2000D\u001b[KitkImageFileReaderIF3: 1.000000\u001b[2000D\u001b[K"
     ]
    }
   ],
   "source": [
    "ImageType = itk.Image[PixelType, Dimension]\n",
    "ImageReaderType = itk.ImageFileReader[ImageType]\n",
    "\n",
    "imageReader1 = ImageReaderType.New()\n",
    "imageReader1.SetFileName(CTImageFileName)\n",
    "imageReader1.Update()\n",
    "\n",
    "CTImage = imageReader1.GetOutput()\n",
    "\n",
    "imageReader2 = ImageReaderType.New()\n",
    "imageReader2.SetFileName(ExpertMaskImageFileName)\n",
    "imageReader2.Update()\n",
    "\n",
    "ExpertMaskImage = imageReader2.GetOutput()\n",
    "\n",
    "imageReader3 = ImageReaderType.New()\n",
    "imageReader3.SetFileName(MaskImageFileName)\n",
    "imageReader3.Update()\n",
    "\n",
    "MaskImage = imageReader3.GetOutput()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 1: Crop the input volumes to make them of same size as MaskImage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "c:\\src\\tubetk_python_itk\\itk\\modules\\io\\imagebase\\include\\itkImageFileWriter.hxx:290:\nitk::ERROR: ImageFileWriter(000000001700C090): Largest possible region does not fully contain requested paste IO regionPaste IO region: ImageIORegion (000000000021CA00)\n  Dimension: 3\n  Index: 0 0 0 \n  Size: 0 0 0 \nLargest possible region: ImageRegion (000000000021CAC8)\n  Dimension: 3\n  Index: [0, 0, 0]\n  Size: [0, 0, 0]\n",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-27-6049da1a2d19>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     31\u001b[0m \u001b[0mimageWriter\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mSetInput\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcroppedMaskImage\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     32\u001b[0m \u001b[0mimageWriter\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mSetFileName\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mMaskImageFileName\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 33\u001b[1;33m \u001b[0mimageWriter\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mUpdate\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m: c:\\src\\tubetk_python_itk\\itk\\modules\\io\\imagebase\\include\\itkImageFileWriter.hxx:290:\nitk::ERROR: ImageFileWriter(000000001700C090): Largest possible region does not fully contain requested paste IO regionPaste IO region: ImageIORegion (000000000021CA00)\n  Dimension: 3\n  Index: 0 0 0 \n  Size: 0 0 0 \nLargest possible region: ImageRegion (000000000021CAC8)\n  Dimension: 3\n  Index: [0, 0, 0]\n  Size: [0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "boundary = itk.Index[3]()\n",
    "boundary.Fill(10)\n",
    "\n",
    "#Create the crop image filter\n",
    "CropImageFilterType = itktube.CropImage[ImageType, ImageType]\n",
    "cropImageFilter = CropImageFilterType.New()\n",
    "cropImageFilter.SetBoundary(boundary)\n",
    "#cropImageFilter.SetMatchVolume(MaskImage)  #Giving error\n",
    "\n",
    "#Crop Input CTA\n",
    "cropImageFilter.SetInput(CTImage)\n",
    "cropImageFilter.Update()\n",
    "\n",
    "croppedCTImage = cropImageFilter.GetOutput()\n",
    "\n",
    "#Crop Expert Mask\n",
    "cropImageFilter.SetInput(ExpertMaskImage)\n",
    "cropImageFilter.Update()\n",
    "\n",
    "croppedExpertMaskImage = cropImageFilter.GetOutput()\n",
    "\n",
    "#Crop Mask\n",
    "cropImageFilter.SetInput(MaskImage)\n",
    "cropImageFilter.Update()\n",
    "\n",
    "croppedMaskImage = cropImageFilter.GetOutput()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 2: Resample the cropped images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "boundary = itk.Index[3]()\n",
    "boundary.Fill(10)\n",
    "\n",
    "#Create the resample image filter\n",
    "ResampleImageFilterType = itktube.ResampleImage[ImageType, ImageType]\n",
    "\n",
    "#Resample Input CTA\n",
    "resampleImageFilter1 = ResampleImageFilterType.New()\n",
    "resampleImageFilter1.SetInput(croppedCTImage)\n",
    "resampleImageFilter1.SetMakeIsotropic(True)\n",
    "resampleImageFilter1.SetInterpolator(\"Sinc\")\n",
    "\n",
    "resampleCTImage = resampleImageFilter1.GetOutput()\n",
    "\n",
    "#Resample Expert Mask\n",
    "resampleImageFilter2 = ResampleImageFilterType.New()\n",
    "resampleImageFilter2.SetInput(croppedExpertMaskImage)\n",
    "resampleImageFilter2.SetMakeIsotropic(True)\n",
    "resampleImageFilter2.SetInterpolator(\"NearestNeighbor\")\n",
    "\n",
    "resampleExpertMaskImage = resampleImageFilter2.GetOutput()\n",
    "\n",
    "#Resample Mask\n",
    "resampleImageFilter3 = ResampleImageFilterType.New()\n",
    "resampleImageFilter3.SetInput(croppedMaskImage)\n",
    "resampleImageFilter3.SetMakeIsotropic(True)\n",
    "resampleImageFilter3.SetInterpolator(\"NearestNeighbor\")\n",
    "\n",
    "resampleMaskImage = resampleImageFilter3.GetOutput()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 3: Create Mask-only images. this step required Image Math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# create resampleMaskImage -> erodedResampleMaskImage\n",
    "# resampleExpertMaskImage -> erodedResampleExpertMaskImage\n",
    "# resampleCTImage -> maskedResampleCTImage\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 4: Compute Training Mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Create image to save not-vessel mask.\n",
    "ShortImageType = itk.Image[itk.S, Dimension]\n",
    "notVesselMaskImage = ShortImageType.New()\n",
    "\n",
    "#Create Compute Training Mask Filter\n",
    "ComputeTrainingMaskFilterType = itktube.ComputeTrainingMask[ImageType]\n",
    "\n",
    "computeTrainingMaskFilter = ComputeTrainingMaskFilterType.New()\n",
    "computeTrainingMaskFilter.SetInput(erodedResampleExpertMaskImage)\n",
    "computeTrainingMaskFilter.SetNotVesselMask(notVesselMaskImage)\n",
    "computeTrainingMaskFilter.SetGap(0.5)\n",
    "computeTrainingMaskFilter.SetNotVesselWidth(2)\n",
    "computeTrainingMaskFilter.Update()\n",
    "\n",
    "expertTrainMaskImage = computeTrainingMaskFilter.GetOutput()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 5: Enhance Vessels in maskedResampleCTImage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "DiscriminantInfoFileName = os.path.join(TubeTK_BUILD_DIR, 'Temporary\\\\vascularModel.mrs')\n",
    "enhancedScalesExpertMaskImage = ImageType.New()\n",
    "\n",
    "# Create EnhanceTubesUsingDiscriminantAnalysis Filter\n",
    "EnhanceTubesUsingDiscriminantAnalysisFilterType = itktube.EnhanceTubesUsingDiscriminantAnalysis[ImageType, ImageType]\n",
    "\n",
    "ETUDAFilter = EnhanceTubesUsingDiscriminantAnalysisFilterType.New()\n",
    "ETUDAFilter.SetInput(maskedResampleCTImage)\n",
    "ETUDAFilter.SetLabelMap(expertTrainMaskImage)\n",
    "ETUDAFilter.SetTubeId(255)\n",
    "ETUDAFilter.SetBackgroundId(127)\n",
    "ETUDAFilter.SetSaveDiscriminantInfo(DiscriminantInfoFileName)\n",
    "ETUDAFilter.SetOutputSeedScaleImage(enhancedScalesExpertMaskImage)\n",
    "ETUDAFilter.SetTubeScales(0.4,0.8,1.2,1.6)\n",
    "\n",
    "enhancedExpertMaskImage = ETUDAFilter.GetOutput()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 6: Compute Segment tubes Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "vasculaModelParameterFileName = os.path.join(TubeTK_BUILD_DIR, 'Temporary\\\\vascularModel.mtp')\n",
    "\n",
    "# Create SegmentTubesParameters Filter\n",
    "ComputeSegmentTubesParametersFilterType = itktube.ComputeSegmentTubesParameters[ImageType]\n",
    "\n",
    "CSTPFilter = ComputeSegmentTubesParametersFilterType.New()\n",
    "CSTPFilter.SetInput(maskedResampleCTImage)\n",
    "CSTPFilter.SetMaskImage(expertTrainMaskImage)\n",
    "CSTPFilter.SetScaleImage(enhancedScalesExpertMaskImage)\n",
    "CSTPFilter.SetParametersFileName(vasculaModelParameterFileName)\n",
    "CSTPFilter.Update()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 7: This step requires Image Math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# enhancedExpertMaskImage -> vesselEnhancedExpertMaskImage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 8: Compute Seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Create SegmentBinaryImageSkeleton Filter\n",
    "SegmentBinaryImageSkeletonFilterType = itktube.SegmentBinaryImageSkeleton[Imagetype]\n",
    "\n",
    "SBISFilter = SegmentBinaryImageSkeletonFilterType.New()\n",
    "SBISFilter.SetInput(vesselEnhancedExpertMaskImage)\n",
    "SBISFilter.Update()\n",
    "\n",
    "seedsVesselEnhancedExpertMaskImage = SBISFilter.GetOutput()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STEP 9: Segment Tubes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "outputVesselsFileName = os.path.join(TubeTK_BUILD_DIR, 'Temporary\\\\outputVessels.tre')\n",
    "\n",
    "\n",
    "\n",
    "# Create SegmentTubes Filter\n",
    "SegmentTubesFilterType = itktube.SegmentTubes[ImageType]\n",
    "\n",
    "SegmenttubesFilter = SegmentTubesFilterType.New()\n",
    "SegmenttubesFilter.SetInput(maskedResampleCTImage)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
